#include "mm.h"
#include "system.h"
#include "kernel.h"

struct bucket_desc
{
    void *page;
    struct bucket_desc* next;
    void *freeptr;
    unsigned short refcnt;
    unsigned short bucket_size;
} __attribute__((packed));

struct _bucket_dir 
{
    int size;
    struct bucket_desc *chain;
} __attribute__((packed));

struct _bucket_dir bucket_dir[] = {
    {16, (struct bucket_desc*) NULL},
    {32, (struct bucket_desc*) NULL},
    {64, (struct bucket_desc*) NULL},
    {128, (struct bucket_desc*) NULL},
    {256, (struct bucket_desc*) NULL},
    {512, (struct bucket_desc*) NULL},
    {1024, (struct bucket_desc*) NULL},
    {2048, (struct bucket_desc*) NULL},
    {4096, (struct bucket_desc*) NULL},
    {0, (struct bucket_desc*) NULL},
};

struct bucket_desc* free_bucket_desc = (struct bucket_desc*) 0;

static inline void init_bucket_desc(void)
{
    struct bucket_desc* bdesc, *first;
    int i;

    first = bdesc = (struct bucket_desc*) get_free_page();
    if (!bdesc)
    {
        return NULL;
    }
    
    for (i = PAGE_SIZE / sizeof(struct bucket_desc); i > 1; i--)
    {
        bdesc->next = bdesc + 1;
        bdesc++;
    }
    
    bdesc->next = free_bucket_desc;
    free_bucket_desc = first;
}

void* kmalloc(size_t len)
{
    struct _bucket_dir *bdir;
    struct bucket_desc * bdesc;
    void * retval;

    for (bdir = bucket_dir; bdir->size > 0; bdir++)
    {
        if (bdir->size >= len)
        {
            break;
        }
        
    }

    if (bdir->size == 0)
    {
        printk("malloc called with impossibly large argement (%d)\n", len);
        return NULL;
    }
    
    for (bdesc = bdir->chain; bdesc; bdesc = bdesc->next)
    {
        if (bdesc->freeptr)
        {
            break;
        }
    }

    if (!bdesc)
    {
        char * cp;
        int i;

        if (!free_bucket_desc)
        {
            init_bucket_desc();
        }
        
        bdesc = free_bucket_desc;
        free_bucket_desc = bdesc->next;
        bdesc->refcnt = 0;
        bdesc->bucket_size = bdir->size;
        bdesc->page = bdesc->freeptr = (void*) (cp = (char*) get_free_page());

        if (!cp)
        {
            return NULL;
        }
        
        for (i = PAGE_SIZE / bdir->size; i > 1; i--)
        {
            *((char**) cp) = cp + bdir->size;
            cp += bdir->size;
        }

        *((char**) cp) = 0;
        bdesc->next = bdir->chain;
        bdir->chain = bdesc;
    }
    
    retval = (void*) bdesc->freeptr;
    bdesc->freeptr = *((void**) retval);
    bdesc->refcnt++;
    return retval;
}

void kfree_s(void* obj, int size)
{
    void* page;
    struct _bucket_dir *bdir;
    struct bucket_desc *bdesc, *prev;
    bdesc = prev = 0;

    page = (void*) ((unsigned long) obj & 0xfffff000);
    for (bdir = bucket_dir; bdir->size; bdir++)
    {
        prev = 0;
        if (bdir->size < size)
        {
            continue;
        }
        for (bdesc = bdir->chain; bdesc; bdesc = bdesc->next)
        {
            if (bdesc->page == page)
            {
                goto found;
            }
            prev = bdesc;
        }
        
    }
    return;
found:
    *((void**) obj) = bdesc->freeptr;
    bdesc->freeptr = obj;
    bdesc->refcnt--;
    if (bdesc->refcnt == 0)
    {
        if ((prev && (prev->next != bdesc)) ||
        (!prev && (bdir->chain != bdesc)))
        {
            for (prev = bdir->chain; prev; prev = prev->next)
            {
                if (prev->next == bdesc)
                {
                    break;
                }
            }
        }

        if (prev) 
        {
            prev->next = bdesc->next;
        } 
        else 
        {
            if (bdir->chain != bdesc)
            {
                return;
            }
            bdir->chain = bdesc->next;
        }
        free_page((unsigned long) bdesc->page);
        bdesc->next = free_bucket_desc;
        free_bucket_desc = bdesc;
    }
    return;
}